package com.zmc.dynamic.shard.strategy;

import com.google.common.collect.Range;
import com.zmc.dynamic.shard.rule.RouteRule;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;
import org.springframework.util.CollectionUtils;

import java.util.*;

/**
 * 多分片键策略
 * 注意：crud只有带分片键才会走分片键策略，根据分片键编写规则路由到指定表中
 */
public class ComplexKeyGeneratorStrategy implements ComplexKeysShardingAlgorithm<String> {

    public static final String SHARD_COMPLEX_STRATEGY = "complex_strategy";

    private Properties props = new Properties();

    /**
     * collection: 当前分片键所有表list
     * complexKeysShardingValue: 多对（k,v）值，和单分片键格式相同
     */
    @Override
    public Collection<String> doSharding(Collection<String> collection, ComplexKeysShardingValue<String> complexKeysShardingValue) {
        //分片键逻辑，可根据项目需求编写
        Map<String, Collection<String>> valuesMap = complexKeysShardingValue.getColumnNameAndShardingValuesMap();
        Set<String> tables = new HashSet<>();
        if (!CollectionUtils.isEmpty(valuesMap)){
            tables.addAll(RouteRule.getTablesByShardingValues(collection, valuesMap));
        }
        Map<String, Range<String>> rangeValuesMap = complexKeysShardingValue.getColumnNameAndRangeValuesMap();
        if (!CollectionUtils.isEmpty(rangeValuesMap)){
            tables.addAll(RouteRule.getTablesByRangeValues(collection, rangeValuesMap));
        }
        return CollectionUtils.isEmpty(tables) ? collection : tables;
    }

    @Override
    public Properties getProps() {
        return props;
    }

    @Override
    public void init(Properties properties) {
        this.props = properties;
    }

    @Override
    public String getType() {
        return SHARD_COMPLEX_STRATEGY;
    }
}
